﻿using System;
using System.Runtime.InteropServices;

namespace DetoursNET
{
    public sealed class DetoursHook<TDelegate> : IDisposable where TDelegate : Delegate
    {
        public IntPtr OriginalPtr { get; private set; }

        public IntPtr OriginalFuncPtr { get; private set; }

        public IntPtr DetourPtr { get; private set; }

        public TDelegate OrigFuncDelegate { get; private set; }

        public TDelegate DetourFuncDelegate { get; private set; }

        public DetoursHook(string moduleName, string funcName, Func<DetoursHook<TDelegate>, TDelegate> delegateGetter)
            : this(Helper.GetProcAddress(moduleName, funcName), delegateGetter)
        {

        }

        public DetoursHook(string moduleName, string funcName, TDelegate targetFuncDelegate)
            : this(Helper.GetProcAddress(moduleName, funcName), targetFuncDelegate)
        {

        }

        public DetoursHook(IntPtr originalPtr, Func<DetoursHook<TDelegate>, TDelegate> delegateGetter)
        {
            DetourFuncDelegate = delegateGetter(this);
            DetourPtr = Marshal.GetFunctionPointerForDelegate(DetourFuncDelegate);

            Hook(originalPtr, DetourPtr);
        }

        public DetoursHook(IntPtr originalPtr, TDelegate targetFuncDelegate)
        {
            DetourFuncDelegate = targetFuncDelegate;
            DetourPtr = Marshal.GetFunctionPointerForDelegate(targetFuncDelegate);

            Hook(originalPtr, DetourPtr);
        }

        public DetoursHook(string moduleName, string funcName, IntPtr detourPtr)
            : this(Helper.GetProcAddress(moduleName, funcName), detourPtr)
        {

        }

        public DetoursHook(IntPtr originalPtr, IntPtr detourPtr)
        {
            Hook(originalPtr, detourPtr);
        }

        public void Dispose()
        {
            try
            {
                Native.DetourTransactionBegin();

                Native.DetourUpdateThread(Native.WinApi.GetCurrentThread());

                var origFuncPtr = OriginalFuncPtr;

                Native.DetourDetach(ref origFuncPtr, DetourPtr);

                Native.DetourTransactionCommit();
            }
            catch (Exception)
            {
                Native.DetourTransactionAbort();
                throw;
            }
            finally
            {
                OrigFuncDelegate = default;
                OriginalFuncPtr = OriginalPtr;
            }
        }

        private void Hook(IntPtr originalPtr, IntPtr detourPtr)
        {
            OriginalPtr = originalPtr;
            OriginalFuncPtr = originalPtr;
            DetourPtr = detourPtr;

            try
            {
                Native.DetourTransactionBegin();

                Native.DetourUpdateThread(Native.WinApi.GetCurrentThread());

                var originalFuncPtr = OriginalFuncPtr;

                Native.DetourAttach(ref originalFuncPtr, DetourPtr);

                Native.DetourTransactionCommit();

                OriginalFuncPtr = originalFuncPtr;

                OrigFuncDelegate = Marshal.GetDelegateForFunctionPointer<TDelegate>(OriginalFuncPtr);
            }
            catch (Exception)
            {
                Native.DetourTransactionAbort();
                throw;
            }
        }
    }
}